import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
from picanet_resnet18 import PiCANetResnet18
from dataset import ECSSDDataset
from tqdm import tqdm

def train(args, device):
    image_dir = os.path.join(args.data_path, 'images')
    mask_dir = os.path.join(args.data_path, 'ground_truth_mask')

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
    mask_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])

    dataset = ECSSDDataset(image_dir, mask_dir, transform, mask_transform)

    # Split dataset: 90% train, 10% val
    val_size = int(0.1 * len(dataset))
    train_size = len(dataset) - val_size
    generator = torch.Generator().manual_seed(42)
    train_dataset, val_dataset = random_split(dataset, [train_size, val_size], generator=generator)

    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)

    model = PiCANetResnet18().to(device)

    criterion = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    os.makedirs(args.save_path, exist_ok=True)

    best_mae = float('inf')

    for epoch in range(args.epochs):
        model.train()
        running_loss = 0.0
        scale_stats = []  # ⬅️ 每轮用于收集 scale 权重统计

        loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}")
        for images, masks in loop:
            images = images.to(device)
            masks = masks.to(device)

            # ⬇️ 支持返回 scale_weights
            output = model(images)
            if isinstance(output, tuple):  # (pred, scale_weights)
                pred, scale_weights = output
                avg = scale_weights.mean(dim=[0, 3, 4])  # (3,)
                scale_stats.append(avg.cpu())
            else:
                pred = output

            loss = criterion(pred, masks)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            loop.set_postfix(loss=loss.item())

        print(f"Epoch [{epoch+1}/{args.epochs}], Loss: {running_loss / len(train_loader):.4f}")

        # ⬇️ 打印当前 epoch 的 scale attention 平均分布
        if scale_stats:
            scale_stats = torch.stack(scale_stats, dim=0).mean(dim=0)
            print(f"[Epoch {epoch+1}] Avg Scale Weights → Local: {scale_stats[0]:.3f}, Mid: {scale_stats[1]:.3f}, Global: {scale_stats[2]:.3f}")

        # --- Validation ---
        model.eval()
        val_mae = 0.0
        with torch.no_grad():
            for val_images, val_masks in val_loader:
                val_images = val_images.to(device)
                val_masks = val_masks.to(device)

                val_pred = model(val_images)
                if isinstance(val_pred, tuple):  # 忽略 scale_weights during eval
                    val_pred = val_pred[0]
                val_pred = torch.sigmoid(val_pred)
                val_mae += torch.mean(torch.abs(val_pred - val_masks)).item()

        val_mae /= len(val_loader)
        print(f"Validation MAE: {val_mae:.4f}")

        # --- Save best model ---
        if val_mae < best_mae:
            best_mae = val_mae
            best_model_path = os.path.join(args.save_path, "best_model.pth")
            torch.save(model.state_dict(), best_model_path)
            print(f"\n✅ Saved best model at epoch {epoch+1} with MAE {best_mae:.4f}\n")
